# -*- coding: utf-8 -*-

import os
import math
import pandas as pd
import openpyxl
from math import sqrt
from numpy import concatenate
import matplotlib.pyplot as plt
import numpy as np
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from tensorflow.keras.layers import *
from tensorflow.keras.models import *
from tensorflow.keras import Input, Model
import warnings
from prettytable import PrettyTable
warnings.filterwarnings("ignore")
import tensorflow as tf
from tensorflow.keras.callbacks import EarlyStopping
from matplotlib.font_manager import FontProperties

# 设置中文字体
font = FontProperties(family='SimSun', size=7.5)

# 加载数据
file_path = "滤波后的features新顺序版414611.xlsx"
sheet_names = pd.ExcelFile(file_path).sheet_names  # 获取所有 sheet 名称

# 加载每个 sheet 对应的电池数据，并保留电池名称
battery_data = {}
battery_names = []
for sheet in sheet_names:
    df = pd.read_excel(file_path, sheet_name=sheet)
    df['battery_name'] = sheet  # 为每个电池数据添加名称列
    battery_data[sheet] = df
    battery_names.extend([sheet] * len(df))  # 扩展电池名称列表

# 将数据拼接到一起
all_data = pd.concat(battery_data.values(), ignore_index=True)

# 将数据转换为 NumPy 数组，去掉电池名称列
values = all_data.drop(columns=['battery_name']).values

# 创建电池名称对应的 NumPy 数组
battery_names = np.array(battery_names)

# 定义三种特征选择方式
feature_sets = {
    "模型Ⅰ": values[:, :8],  # 前8列作为特征
    "模型Ⅱ": values[:, 8:11],  # 9、10、11列作为特征
    "模型Ⅲ": values[:, :11]  # 前11列作为特征
}

# 目标列
Y = values[:, -1]  # 最后一列为目标

# 分割训练集、验证集和测试集：10个训练集，2个验证集，2个测试集
train_sheets = sheet_names[:10]  # 前10个sheet作为训练集
val_sheets = sheet_names[10:12]  # 第11和12个sheet作为验证集
test_sheets = sheet_names[12:14]  # 第13和14个sheet作为测试集

# 提取训练集、验证集和测试集的索引
train_indices = []
val_indices = []
test_indices = []
current_idx = 0
for sheet in sheet_names:
    sheet_len = len(battery_data[sheet])
    if sheet in train_sheets:
        train_indices.extend(range(current_idx, current_idx + sheet_len))
    elif sheet in val_sheets:
        val_indices.extend(range(current_idx, current_idx + sheet_len))
    else:
        test_indices.extend(range(current_idx, current_idx + sheet_len))
    current_idx += sheet_len

# 分割数据
X_train = {name: features[train_indices] for name, features in feature_sets.items()}
Y_train = Y[train_indices]
X_val = {name: features[val_indices] for name, features in feature_sets.items()}
Y_val = Y[val_indices]
X_test = {name: features[test_indices] for name, features in feature_sets.items()}
Y_test = Y[test_indices]
shuffled_battery_names_test = battery_names[test_indices]  # 测试集的电池名称
shuffled_original_indices_test = np.array(test_indices)    # 测试集的原始索引
shuffled_battery_names_val = battery_names[val_indices]    # 验证集的电池名称
shuffled_original_indices_val = np.array(val_indices)      # 验证集的原始索引

# 定义注意力机制层
def attention_layer(inputs, single_attention_vector=False):
    time_steps = tf.keras.backend.int_shape(inputs)[1]
    input_dim = tf.keras.backend.int_shape(inputs)[2]
    a = Permute((2, 1))(inputs)
    a = Reshape((input_dim, time_steps))(a)
    a = Dense(time_steps, activation='softmax')(a)
    if single_attention_vector:
        a = Lambda(lambda x: tf.keras.backend.mean(x, axis=1), name='dim_reduction')(a)
        a = RepeatVector(input_dim)(a)
    a_probs = Permute((2, 1), name='attention_vec')(a)
    output_attention_mul = Multiply()([inputs, a_probs])
    return output_attention_mul

# 定义LSTM+Attention模型
def lstm_attention_model(input_shape):
    inputs = Input(shape=(input_shape[1], input_shape[2]))
    lstm_out = LSTM(128, return_sequences=True)(inputs)
    attention_mul = attention_layer(lstm_out)
    attention_flatten = Flatten()(attention_mul)
    outputs = Dense(1)(attention_flatten)
    model = Model(inputs=[inputs], outputs=[outputs])
    model.compile(loss='mse', optimizer='Adam')
    model.summary()
    return model

# 定义MAPE计算函数
def mape(y_true, y_pred):
    record = []
    for index in range(len(y_true)):
        temp_mape = np.abs((y_pred[index] - y_true[index]) / y_true[index])
        record.append(temp_mape)
    return np.mean(record) * 100

# 定义评估函数
def evaluate_forecasts(Ytest, predicted_data, n_out):
    mse_dic = []
    rmse_dic = []
    mae_dic = []
    mape_dic = []
    r2_dic = []
    table = PrettyTable(['测试集指标', 'MSE', 'RMSE', 'MAE', 'MAPE', 'R2'])
    for i in range(n_out):
        actual = [float(row[i]) for row in Ytest]
        predicted = [float(row[i]) for row in predicted_data]
        mse = mean_squared_error(actual, predicted)
        mse_dic.append(mse)
        rmse = sqrt(mean_squared_error(actual, predicted))
        rmse_dic.append(rmse)
        mae = mean_absolute_error(actual, predicted)
        mae_dic.append(mae)
        MApe = mape(actual, predicted)
        mape_dic.append(MApe)
        r2 = r2_score(actual, predicted)
        r2_dic.append(r2)
        if n_out == 1:
            strr = '预测结果指标：'
        else:
            strr = '第' + str(i + 1) + '步预测结果指标：'
        table.add_row([strr, mse, rmse, mae, str(MApe) + '%', str(r2 * 100) + '%'])
    print(table)  # 在控制台打印误差表
    return mse_dic, rmse_dic, mae_dic, mape_dic, r2_dic, table

# 数据隔离类，确保训练集、验证集和测试集不混用
class DataIsolation:
    def __init__(self, X_train, Y_train, X_val, Y_val, X_test, Y_test):
        self.X_train = X_train
        self.Y_train = Y_train
        self.X_val = X_val
        self.Y_val = Y_val
        self.X_test = X_test
        self.Y_test = Y_test
        
        # 初始化归一化器
        self.input_scalers = {}
        self.output_scaler = MinMaxScaler()
        
        # 只使用训练集拟合归一化器
        for feature_name in self.X_train:
            self.input_scalers[feature_name] = MinMaxScaler()
            self.X_train[feature_name] = self.input_scalers[feature_name].fit_transform(self.X_train[feature_name])
            self.X_val[feature_name] = self.input_scalers[feature_name].transform(self.X_val[feature_name])
            self.X_test[feature_name] = self.input_scalers[feature_name].transform(self.X_test[feature_name])
        
        self.Y_train = self.output_scaler.fit_transform(self.Y_train.reshape(-1, 1))
        self.Y_val = self.output_scaler.transform(self.Y_val.reshape(-1, 1))
        self.Y_test = self.output_scaler.transform(self.Y_test.reshape(-1, 1))
        
    def get_train_data(self, feature_name):
        X = self.X_train[feature_name].reshape(-1, 1, self.X_train[feature_name].shape[1])
        return X, self.Y_train
    
    def get_val_data(self, feature_name):
        X = self.X_val[feature_name].reshape(-1, 1, self.X_val[feature_name].shape[1])
        return X, self.Y_val
    
    def get_test_data(self, feature_name):
        X = self.X_test[feature_name].reshape(-1, 1, self.X_test[feature_name].shape[1])
        return X, self.Y_test
    
    def inverse_transform(self, y):
        return self.output_scaler.inverse_transform(y)

# 存储结果
results = {}
loss_histories = {}
mape_results = {}
all_errors = []

# 初始化数据隔离
data_iso = DataIsolation(X_train, Y_train, X_val, Y_val, X_test, Y_test)

# 对每种特征选择方式进行训练和预测
for feature_name in feature_sets:
    # 提取训练、验证和测试数据
    X_train_data, Y_train_data = data_iso.get_train_data(feature_name)
    X_val_data, Y_val_data = data_iso.get_val_data(feature_name)
    X_test_data, Y_test_data = data_iso.get_test_data(feature_name)
    
    # 训练模型，使用验证集监控性能
    model = lstm_attention_model(X_train_data.shape)
    early_stop = EarlyStopping(monitor='val_loss', patience=50, restore_best_weights=True)
    history = model.fit(
        X_train_data, Y_train_data,
        validation_data=(X_val_data, Y_val_data),
        batch_size=72,
        epochs=1000,
        callbacks=[early_stop],
        verbose=2
    )
    loss_histories[feature_name] = history

    # 在验证集上进行预测
    yhat_val = model.predict(X_val_data)
    predicted_val = data_iso.inverse_transform(yhat_val)
    actual_val = data_iso.inverse_transform(Y_val_data)

    # 在测试集上进行预测
    yhat_test = model.predict(X_test_data)
    predicted_test = data_iso.inverse_transform(yhat_test)
    actual_test = data_iso.inverse_transform(Y_test_data)

    # 存储结果
    results[feature_name] = {
        'val': (actual_val, predicted_val),
        'test': (actual_test, predicted_test)
    }

    # 计算验证集的MAPE
    mape_values_val = []
    for i in range(len(actual_val)):
        mape_values_val.append(np.abs((predicted_val[i] - actual_val[i]) / actual_val[i]) * 100)
    mape_results[feature_name + '_val'] = mape_values_val

    # 计算测试集的MAPE
    mape_values_test = []
    for i in range(len(actual_test)):
        mape_values_test.append(np.abs((predicted_test[i] - actual_test[i]) / actual_test[i]) * 100)
    mape_results[feature_name + '_test'] = mape_values_test
    
    # 评估验证集
    mse_val, rmse_val, mae_val, mape_val_val, r2_val, table_val = evaluate_forecasts(actual_val, predicted_val, 1)
    all_errors.append([feature_name + '_val', mse_val[0], rmse_val[0], mae_val[0], mape_val_val[0], r2_val[0]])

    # 评估测试集
    mse_test, rmse_test, mae_test, mape_val_test, r2_test, table_test = evaluate_forecasts(actual_test, predicted_test, 1)
    all_errors.append([feature_name + '_test', mse_test[0], rmse_test[0], mae_test[0], mape_val_test[0], r2_test[0]])

# 绘制损失曲线
plt.figure(figsize=(3.03, 1.5), dpi=300)
for feature_name, history in loss_histories.items():
    plt.plot(history.history['loss'], label=f'{feature_name}训练集')
    plt.plot(history.history['val_loss'], label=f'{feature_name}验证集')

plt.xlabel('训练轮次', fontproperties=font)
plt.ylabel('损失值', fontproperties=font)
plt.legend(
    prop=font, 
    fontsize=7.5,
    loc='upper right',
    frameon=True, 
    borderpad=0.5,
    labelspacing=0.5,
    borderaxespad=0.5,
    ncol=2
)
plt.tight_layout(pad=0.2)
plt.subplots_adjust(left=0.01, right=0.99, top=0.99, bottom=0.01)
plt.show()

# 绘制每个电池的预测结果与真实值的对比（验证集和测试集）
fig, axes = plt.subplots(4, 1, figsize=(6.03, 12))  # 修改为4行1列，2个验证集 + 2个测试集
fig.subplots_adjust(hspace=0.5)
colors = ['#2ca02c', '#1f77b4', '#d62728']

# 绘制验证集的电池
for idx, battery_name in enumerate(val_sheets):
    battery_indices = np.where(shuffled_battery_names_val == battery_name)[0]
    if len(battery_indices) == 0:
        continue
    ax = axes[idx]
    for feature_idx, (feature_name, result_dict) in enumerate(results.items()):
        Yval, predicted_data = result_dict['val']
        battery_Yval = Yval[battery_indices]
        battery_predicted_data = predicted_data[battery_indices]
        original_indices = shuffled_original_indices_val[battery_indices]
        sorted_indices = np.argsort(original_indices)
        ax.plot(range(1, len(battery_indices)+1), 
                battery_predicted_data[sorted_indices], 
                linestyle="--", 
                linewidth=1,
                label=f'{feature_name}预测值',
                marker="o", 
                markersize=2,
                color=colors[feature_idx])
    battery_Yval = Yval[battery_indices][sorted_indices]
    ax.plot(range(1, len(battery_indices)+1), 
            battery_Yval,
            linestyle="-",
            linewidth=1.5,
            label='实际值',
            marker="x",
            markersize=3,
            color='#ff7f0e')
    ax.set_title(f"{battery_name}容量预测结果 (验证集)", fontproperties=font)
    ax.set_xlabel("循环周期", fontproperties=font)
    ax.set_ylabel("容量 (Ah)", fontproperties=font)
    ax.grid(True, alpha=0.3)
    ax.legend(prop=font, loc='best', frameon=True, shadow=True)

# 绘制测试集的电池
for idx, battery_name in enumerate(test_sheets):
    battery_indices = np.where(shuffled_battery_names_test == battery_name)[0]
    if len(battery_indices) == 0:
        continue
    ax = axes[idx + 2]  # 从第3个子图开始
    for feature_idx, (feature_name, result_dict) in enumerate(results.items()):
        Ytest, predicted_data = result_dict['test']
        battery_Ytest = Ytest[battery_indices]
        battery_predicted_data = predicted_data[battery_indices]
        original_indices = shuffled_original_indices_test[battery_indices]
        sorted_indices = np.argsort(original_indices)
        ax.plot(range(1, len(battery_indices)+1), 
                battery_predicted_data[sorted_indices], 
                linestyle="--", 
                linewidth=1,
                label=f'{feature_name}预测值',
                marker="o", 
                markersize=2,
                color=colors[feature_idx])
    battery_Ytest = Ytest[battery_indices][sorted_indices]
    ax.plot(range(1, len(battery_indices)+1), 
            battery_Ytest,
            linestyle="-",
            linewidth=1.5,
            label='实际值',
            marker="x",
            markersize=3,
            color='#ff7f0e')
    ax.set_title(f"{battery_name}容量预测结果 (测试集)", fontproperties=font)
    ax.set_xlabel("循环周期", fontproperties=font)
    ax.set_ylabel("容量 (Ah)", fontproperties=font)
    ax.grid(True, alpha=0.3)
    ax.legend(prop=font, loc='best', frameon=True, shadow=True)

plt.tight_layout()
plt.savefig('prediction_comparison_vertical.png', dpi=300, bbox_inches='tight')
plt.show()

# 保存预测结果到Excel（验证集和测试集）
with pd.ExcelWriter('预测结果对比461411.xlsx') as writer:
    for battery_name in val_sheets + test_sheets:
        data_dict = {'实际值': []}
        if battery_name in val_sheets:
            battery_indices = np.where(shuffled_battery_names_val == battery_name)[0]
            original_indices = shuffled_original_indices_val[battery_indices]
            sorted_indices = np.argsort(original_indices)
            actual_values = Y_val[battery_indices][sorted_indices].flatten()
            data_dict['实际值'] = actual_values
            for feature_name in feature_sets.keys():
                pred_values = results[feature_name]['val'][1][battery_indices][sorted_indices].flatten()
                data_dict[f'{feature_name}预测值'] = pred_values
            sheet_prefix = "Val_"
        else:
            battery_indices = np.where(shuffled_battery_names_test == battery_name)[0]
            original_indices = shuffled_original_indices_test[battery_indices]
            sorted_indices = np.argsort(original_indices)
            actual_values = Y_test[battery_indices][sorted_indices].flatten()
            data_dict['实际值'] = actual_values
            for feature_name in feature_sets.keys():
                pred_values = results[feature_name]['test'][1][battery_indices][sorted_indices].flatten()
                data_dict[f'{feature_name}预测值'] = pred_values
            sheet_prefix = "Test_"
        df = pd.DataFrame(data_dict)
        sheet_name = f"{sheet_prefix}{battery_name[:27]}" if len(battery_name) > 27 else f"{sheet_prefix}{battery_name}"
        df.to_excel(writer, sheet_name=sheet_name, index=False)
    info_df = pd.DataFrame({
        '说明': [
            '本文件包含各电池容量预测结果',
            'Val_前缀表示验证集电池，Test_前缀表示测试集电池',
            '实际值：电池实际容量测量值',
            '模型Ⅰ预测值：使用前8个特征的预测结果',
            '模型Ⅱ预测值：使用9-11特征的预测结果',
            '模型Ⅲ预测值：使用全部11个特征的预测结果'
        ]
    })
    info_df.to_excel(writer, sheet_name='说明', index=False)

# 计算每个电池在不同模型下的平均MAPE（验证集和测试集）
battery_mape_means = {}
for battery_name in val_sheets + test_sheets:
    if battery_name in val_sheets:
        battery_indices = np.where(shuffled_battery_names_val == battery_name)[0]
        mape_key = '_val'
    else:
        battery_indices = np.where(shuffled_battery_names_test == battery_name)[0]
        mape_key = '_test'
    if len(battery_indices) == 0:
        continue
    mean_mape = []
    for feature_name in feature_sets.keys():
        mape_values = np.array(mape_results[feature_name + mape_key])[battery_indices]
        mean_mape.append(np.mean(mape_values))
    battery_mape_means[battery_name] = mean_mape

# 提取纯数字编号
def extract_battery_number(name):
    return ''.join(filter(str.isdigit, name))

cleaned_battery_names = [extract_battery_number(name) for name in battery_mape_means.keys()]
model1_means = [v[0] for v in battery_mape_means.values()]
model2_means = [v[1] for v in battery_mape_means.values()]
model3_means = [v[2] for v in battery_mape_means.values()]

# 创建堆叠柱形图
fig, ax = plt.subplots(figsize=(2.4, 1.4), dpi=300)
colors = ['#1f77b4', '#ff7f0e', '#2ca02c']
p1 = ax.bar(cleaned_battery_names, model1_means, color=colors[0], label='模型Ⅰ')
p2 = ax.bar(cleaned_battery_names, model2_means, bottom=model1_means, color=colors[1], label='模型Ⅱ')
p3 = ax.bar(cleaned_battery_names, model3_means, 
           bottom=np.array(model1_means) + np.array(model2_means), 
           color=colors[2], label='模型Ⅲ')
ax.set_xlabel('电池编号', fontproperties=font)
ax.set_ylabel('平均绝对百分比误差 (%)', fontproperties=font)
plt.xticks(
    ha='right', 
    fontproperties=font,
    ticks=range(len(cleaned_battery_names)),
    labels=cleaned_battery_names
)
plt.yticks(fontproperties=font)
ax.legend(loc='upper left', prop=font, frameon=True, borderpad=0.5, labelspacing=0.5)
plt.tight_layout(pad=0.2)
plt.subplots_adjust(left=0.01, right=0.99, top=0.99, bottom=0.01)
plt.savefig('stacked_bar_chart.png', dpi=300, bbox_inches='tight')
plt.show()

# 为每个电池的每种预测生成误差表（验证集和测试集）
for battery_name in val_sheets + test_sheets:
    if battery_name in val_sheets:
        battery_indices = np.where(shuffled_battery_names_val == battery_name)[0]
        result_key = 'val'
    else:
        battery_indices = np.where(shuffled_battery_names_test == battery_name)[0]
        result_key = 'test'
    if len(battery_indices) == 0:
        continue
    for feature_name, result_dict in results.items():
        Yset, predicted_data = result_dict[result_key]
        battery_Yset = Yset[battery_indices]
        battery_predicted_data = predicted_data[battery_indices]
        mse = mean_squared_error(battery_Yset, battery_predicted_data)
        rmse = sqrt(mse)
        mae = mean_absolute_error(battery_Yset, battery_predicted_data)
        mape_val = mape(battery_Yset, battery_predicted_data)
        r2 = r2_score(battery_Yset, battery_predicted_data)
        set_type = "验证集" if result_key == 'val' else "测试集"
        table = PrettyTable(['电池名称', '特征选择', '集合类型', 'MSE', 'RMSE', 'MAE', 'MAPE', 'R2'])
        table.add_row([battery_name, feature_name, set_type, mse, rmse, mae, f"{mape_val:.2f}%", f"{r2 * 100:.2f}%"])
        print(table)

# 输出所有特征选择的误差总表
error_table = PrettyTable(['特征选择', '集合类型', 'MSE', 'RMSE', 'MAE', 'MAPE', 'R2'])
for error in all_errors:
    feature_name, mse, rmse, mae, mape_val, r2 = error
    set_type = "验证集" if '_val' in feature_name else "测试集"
    feature_name = feature_name.replace('_val', '').replace('_test', '')
    error_table.add_row([feature_name, set_type, mse, rmse, mae, f"{mape_val:.2f}%", f"{r2 * 100:.2f}%"])
print("所有特征选择的误差总表：")
print(error_table)